import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import ks_2samp

# 生成增强型模拟数据
def generate_population(size=2000):
    np.random.seed(42)
    return pd.DataFrame({
        'id': range(size),
        'age': np.random.normal(35, 10, size).astype(int),
        'income': np.random.lognormal(8.5, 0.8, size),
        'region': np.random.choice(['Urban','Rural'], size, p=[0.7,0.3]),
        'education': np.random.choice(['High School','Bachelor','Master','PhD'], 
                                  size, p=[0.35,0.4,0.2,0.05]),
        'cluster': np.random.randint(1, 51, size)
    })

# 初始化数据集
data = generate_population()

# 概率抽样方法实现
## 1. 多阶段抽样（省-市-街道）
def multistage_sampling(df):
    # 第一阶段：抽取10个群集
    clusters = df['cluster'].unique()
    stage1 = np.random.choice(clusters, 10, replace=False)
    
    # 第二阶段：从每个群集抽取5%个体
    return df[df['cluster'].isin(stage1)].groupby('cluster', group_keys=False)\
           .apply(lambda x: x.sample(frac=0.05))

## 2. 分层比例抽样（按地区和教育程度）
def stratified_proportional(df):
    strata = df.groupby(['region', 'education']).size()
    sample_size = 300
    return df.groupby(['region', 'education'], group_keys=False)\
           .apply(lambda x: x.sample(n=int(sample_size*len(x)/len(df))))

# 非概率抽样方法实现
## 3. 滚雪球抽样（模拟社交网络）
def snowball_sampling(df):
    seed = df.sample(10)
    referrals = df[df['id'].isin(seed['id']//100*100 + np.random.randint(1,100,50))]
    return pd.concat([seed, referrals]).drop_duplicates()

# 评估模块增强版
def evaluate_samples(population, sample, method_name):
    metrics = {
        '方法': method_name,
        '样本量': len(sample),
        '年龄均值差异': (population['age'].mean() - sample['age'].mean()).round(2),
        '收入中位数差异': (population['income'].median() - sample['income'].median()).round(2),
        '地区分布KS统计量': ks_2samp(population['region'], sample['region'])[0].round(3)
    }
    
    # 可视化分析
    fig, ax = plt.subplots(1,2, figsize=(12,5))
    for col, ax in zip(['age','income'], ax):
        sns.kdeplot(population[col], ax=ax, label='总体')
        sns.kdeplot(sample[col], ax=ax, label='样本')
        ax.set_title(f'{col}分布对比')
    plt.tight_layout()
    return pd.DataFrame([metrics])

# 执行所有抽样评估
results = pd.concat([
    evaluate_samples(data, multistage_sampling(data), '多阶段抽样'),
    evaluate_samples(data, stratified_proportional(data), '分层比例抽样'),
    evaluate_samples(data, snowball_sampling(data), '滚雪球抽样')
])

print(results.to_markdown(index=False))
